package cn.iocoder.yudao.module.data.util;

import cn.hutool.core.util.StrUtil;
import cn.hutool.db.DbUtil;
import cn.hutool.db.ds.DSFactory;
import cn.iocoder.yudao.module.data.service.config.ParserConfig;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.enums.DataFormatTypeEnum;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.enums.EngineType;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.s2sql.LLMReq;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.s2sql.Text2SQLExemplar;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

import static cn.iocoder.yudao.module.data.service.config.ParserConfig.*;


@Slf4j
public class PromptHelper  {

    public List<List<Text2SQLExemplar>> getFewShotExemplars(LLMReq llmReq) {
       ParserConfig parserConfig = new ParserConfig();
        int exemplarRecallNumber =
                Integer.parseInt(parserConfig.getParameterValue(PARSER_EXEMPLAR_RECALL_NUMBER));
        int fewShotNumber =
                Integer.parseInt(parserConfig.getParameterValue(PARSER_FEW_SHOT_NUMBER));
        int selfConsistencyNumber =
                Integer.parseInt(parserConfig.getParameterValue(PARSER_SELF_CONSISTENCY_NUMBER));

        List<Text2SQLExemplar> exemplars = Lists.newArrayList();
        exemplars.addAll(llmReq.getDynamicExemplars());

        int recallSize = exemplarRecallNumber - llmReq.getDynamicExemplars().size();
//        if (recallSize > 0) {
//            exemplars.addAll(exemplarService.recallExemplars(llmReq.getQueryText(), recallSize));
//        }

        List<List<Text2SQLExemplar>> results = new ArrayList<>();
        // use random collection of exemplars for each self-consistency inference
        for (int i = 0; i < selfConsistencyNumber; i++) {
            List<Text2SQLExemplar> shuffledList = new ArrayList<>(exemplars);
            Collections.shuffle(shuffledList);
            results.add(shuffledList.subList(0, Math.min(shuffledList.size(), fewShotNumber)));
        }

        return results;
    }

    public String buildSideInformation(LLMReq llmReq) {
        String currentDate = llmReq.getCurrentDate();
        List<String> sideInfos = Lists.newArrayList();
        sideInfos.add(String.format("CurrentDate=[%s]", currentDate));

        if (StringUtils.isNotEmpty(llmReq.getPriorExts())) {
            sideInfos.add(String.format("PriorKnowledge=[%s]", llmReq.getPriorExts()));
        }

        LLMReq.LLMSchema schema = llmReq.getSchema();
        if (!isSupportWith(schema.getDatabaseType(), schema.getDatabaseVersion())) {
            sideInfos.add("[Database does not support with statement]");
        }

        String termStr = buildTermStr(llmReq);
        if (StringUtils.isNotEmpty(termStr)) {
            sideInfos.add(String.format("DomainTerms=[%s]", termStr));
        }

        return String.join(",", sideInfos);
    }

    /**
     *
     * @param llmReq
     * @return
     */
    public String buildSchemaWithDDL(LLMReq llmReq) {
        //通过jdbc获取表的DDL
//        String tableStr = llmReq.getSchema().get();
        String tableStr = "bi_indicator";
        String databaseType = llmReq.getSchema().getDatabaseType();
        String databaseVersion = llmReq.getSchema().getDatabaseVersion();
        String ddl = "";
        if (EngineType.MYSQL.getName().equalsIgnoreCase(databaseType)) {
            ddl = String.format("SHOW CREATE TABLE %s", tableStr);
        } else if (EngineType.CLICKHOUSE.getName().equalsIgnoreCase(databaseType)) {
            ddl = String.format("SHOW CREATE TABLE %s.%s", databaseVersion, tableStr);
        } else {
            ddl = String.format("SHOW CREATE TABLE %s", tableStr);
        }
        // 通过jdbc获取表的DDL
        Connection connection = null;
        try {
            String jdbcUrl = llmReq.getSchema().getJdbcUrl();
            String userName = llmReq.getSchema().getUserName();
            String password = llmReq.getSchema().getPassword();

            connection = DriverManager.getConnection(jdbcUrl, userName, password);
            if (connection != null) {
                // 执行DDL语句
                ResultSet resultSet = connection.createStatement().executeQuery(ddl);
                StringBuilder ddlResult = new StringBuilder();

                // 处理DDL结果
                // 这里假设DDL结果是一个字符串，你可以根据实际情况进行处理
                // 例如，提取表名、字段名等信息
                // 这里仅仅是一个示例，你可以根据实际需求进行处理
                // 生成DDL字符串
                while (resultSet.next()) {
                    String ddlString = resultSet.getString(2);
                    ddlResult.append(ddlString);
                }

                return ddlResult.toString();
            } else {
                log.error("Connection is null");
            }
        } catch (Exception e) {
            log.error("Error executing DDL: {}", e.getMessage());
        } finally {
            if (connection != null) {
                try {
                    connection.close();
                } catch (Exception e) {
                    log.error("Error closing connection: {}", e.getMessage());
                }
            }
        }

        return "";

    }

    /**
     *
     * @param llmReq
     * @return
     */
    public String buildSchemaStr(LLMReq llmReq) {
        String tableStr = llmReq.getSchema().getDataSetName();

        List<String> metrics = Lists.newArrayList();
        llmReq.getSchema().getMetrics().forEach(metric -> {
            StringBuilder metricStr = new StringBuilder();
            metricStr.append("<");
            metricStr.append(metric.getName());
            if (!CollectionUtils.isEmpty(metric.getAlias())) {
                StringBuilder alias = new StringBuilder();
                metric.getAlias().forEach(a -> alias.append(a).append(","));
                metricStr.append(" ALIAS '").append(alias).append("'");
            }
            if (StringUtils.isNotEmpty(metric.getDataFormatType())) {
                String dataFormatType = metric.getDataFormatType();
                if (DataFormatTypeEnum.DECIMAL.getName().equalsIgnoreCase(dataFormatType)
                        || DataFormatTypeEnum.PERCENT.getName().equalsIgnoreCase(dataFormatType)) {
                    metricStr.append(" FORMAT '").append(dataFormatType).append("'");
                }
            }
            if (StringUtils.isNotEmpty(metric.getDescription())) {
                metricStr.append(" COMMENT '").append(metric.getDescription()).append("'");
            }
            if (StringUtils.isNotEmpty(metric.getDefaultAgg())) {
                metricStr.append(" AGGREGATE '").append(metric.getDefaultAgg().toUpperCase())
                        .append("'");
            }
            metricStr.append(">");
            metrics.add(metricStr.toString());
        });

        List<String> dimensions = Lists.newArrayList();
        llmReq.getSchema().getDimensions().forEach(dimension -> {
            StringBuilder dimensionStr = new StringBuilder();
            dimensionStr.append("<");
            dimensionStr.append(dimension.getName());
            if (!CollectionUtils.isEmpty(dimension.getAlias())) {
                StringBuilder alias = new StringBuilder();
                dimension.getAlias().forEach(a -> alias.append(a).append(";"));
                dimensionStr.append(" ALIAS '").append(alias).append("'");
            }
            if (StringUtils.isNotEmpty(dimension.getTimeFormat())) {
                dimensionStr.append(" FORMAT '").append(dimension.getTimeFormat()).append("'");
            }
            if (StringUtils.isNotEmpty(dimension.getDescription())) {
                dimensionStr.append(" COMMENT '").append(dimension.getDescription()).append("'");
            }
            dimensionStr.append(">");
            dimensions.add(dimensionStr.toString());
        });

        List<String> values = Lists.newArrayList();
        List<LLMReq.ElementValue> elementValueList = llmReq.getSchema().getValues();
        if (elementValueList != null) {
            elementValueList.forEach(value -> {
                StringBuilder valueStr = new StringBuilder();
                String fieldName = value.getFieldName();
                String fieldValue = value.getFieldValue();
                valueStr.append(String.format("<%s='%s'>", fieldName, fieldValue));
                values.add(valueStr.toString());
            });
        }

        String partitionTimeStr = "";
        if (llmReq.getSchema().getPartitionTime() != null) {
            partitionTimeStr =
                    String.format("%s FORMAT '%s'", llmReq.getSchema().getPartitionTime().getName(),
                            llmReq.getSchema().getPartitionTime().getTimeFormat());
        }

        String primaryKeyStr = "";
        if (llmReq.getSchema().getPrimaryKey() != null) {
            primaryKeyStr = String.format("%s", llmReq.getSchema().getPrimaryKey().getName());
        }

        String databaseTypeStr = "";
        if (llmReq.getSchema().getDatabaseType() != null) {
            databaseTypeStr = llmReq.getSchema().getDatabaseType();
        }
        String databaseVersionStr = "";
        if (llmReq.getSchema().getDatabaseVersion() != null) {
            databaseVersionStr = llmReq.getSchema().getDatabaseVersion();
        }

        String template =
                "DatabaseType=[%s], DatabaseVersion=[%s], Table=[%s], PartitionTimeField=[%s], PrimaryKeyField=[%s], "
                        + "Metrics=[%s], Dimensions=[%s], Values=[%s]";
        return String.format(template, databaseTypeStr, databaseVersionStr, tableStr,
                partitionTimeStr, primaryKeyStr, String.join(",", metrics),
                String.join(",", dimensions), String.join(",", values));
    }

    private String buildTermStr(LLMReq llmReq) {
        List<LLMReq.Term> terms = llmReq.getTerms();
        List<String> termStr = Lists.newArrayList();
        if(terms !=null && !terms.isEmpty()) {
            terms.forEach(term -> {
                StringBuilder termsDesc = new StringBuilder();
                String description = term.getDescription();
                termsDesc.append(String.format("<%s COMMENT '%s'>", term.getName(), description));
                termStr.add(termsDesc.toString());
            });
        }

        String ret = "";
        if (!termStr.isEmpty()) {
            ret = String.join(",", termStr);
        }

        return ret;
    }

    public static boolean isSupportWith(String type, String version) {
        if (type.equalsIgnoreCase(EngineType.MYSQL.getName()) && Objects.nonNull(version)
                && StrUtil.compareVersion(version, "8.0") < 0) {
            return false;
        }
        if (type.equalsIgnoreCase(EngineType.CLICKHOUSE.getName()) && Objects.nonNull(version)
                && StrUtil.compareVersion(version, "20.4") < 0) {
            return false;
        }
        return true;
    }
}
